import os
import torch
import numpy as np
from PIL import Image
import sys

# Define paths
gt_path = 'D:/deeplearning/EGNet/data/ground_truth_mask'
pred_path = 'D:/deeplearning/EGNet/data/images'

# Create lists of ground truth and predicted image paths
gt_paths = [os.path.join(gt_path, f'{i:04d}.png') for i in range(1, 301)]
pred_paths = [os.path.join(pred_path, f'{i:04d}.png') for i in range(1, 301)]

def prec_rec(y_true, y_pred, beta2):
    eps = sys.float_info.epsilon
    tp = torch.sum(y_true * y_pred)
    all_p_pred = torch.sum(y_pred)
    all_p_true = torch.sum(y_true)
    
    prec = (tp + eps) / (all_p_pred + eps)
    rec = (tp + eps) / (all_p_true + eps)
    
    return prec, rec

overall_mae = 0
total_prec = 0
total_rec = 0

for j in range(len(gt_paths)):
    gt = np.array(Image.open(gt_paths[j]).convert('LA')) / 255
    pred = np.array(Image.open(pred_paths[j]).convert('LA')) / 255 
    mae = np.sum(np.abs(pred - gt)) / (pred.shape[:2][0] * pred.shape[:2][1])
    
    gt_arr = torch.from_numpy(gt).float()
    pred_arr = torch.from_numpy(pred).float()
    threshold = 216 / 255.0  # Convert to a value between 0 and 1
    y_pred = torch.ge(pred_arr, threshold).float()
    y_true = torch.ge(gt_arr, 128 / 255.0).float()  # Convert to a value between 0 and 1
    y_true1 = torch.reshape(y_true, (1, -1))
    y_pred1 = torch.reshape(y_pred, (1, -1))
    
    prec, rec = prec_rec(y_true1, y_pred1, 0.3)

    total_prec += prec
    total_rec += rec
    overall_mae += mae


beta2 = 0.3 
overall_fb = (1 + beta2) * (total_prec * total_rec) / ((beta2 * total_prec + total_rec) * len(gt_paths))
print('overall_mae', overall_mae / len(gt_paths))
print('overall_fb', overall_fb)
